Add AVX 512 support for bg4 predictor - #965
Open
ajanon wants to merge 3 commits into
Open
Conversation
The popcount driver was hardcoded to u128 words. Make it generic over the word type, with the per-byte popcount supplied as a closure, so wider SIMD words can plug in without duplicating the alignment, body, and tail logic. Bound the word type with bytemuck::Pod: Pod's contract (no padding, every bit pattern valid, zeroable) is exactly what the raw aligned word reads and the byte views below require. Move bytemuck to unconditional dependencies, as this module is now needed on all targets, wasm included. SWAR and Neon paths are unchanged and still operate on u128 words.
Add an opt-in add_data_avx512 entry point on x86-64 that runs vpopcntb on 64-byte words. The feature to detect is avx512bitalg: the byte-wise popcount instructions come from AVX512-BITALG, not AVX512-VPOPCNTDQ (which only covers the 32/64-bit lane variants). Fall back to the SWAR method when the CPU lacks it. The histogram scatter is scalar, so wide words would pay per-byte vector extracts. Feed the scatter 16-byte u128 chunks copied into locals instead: the constant trip count lets the compiler fully unroll it and hoist the per-lane histogram row pointers, and a u128 chunk yields its per-byte popcounts with shifts from registers rather than round-tripping through memory. This tightens the word-size requirement from a multiple of 4 to a multiple of 16. bytemuck's avx512_simd feature (Pod impls for __m512i) is enabled for x86_64 builds only. On a Ryzen 9 7950X: SWAR ~3000 MB/s, AVX-512 ~3300-3400 MB/s. The path is throughput-bound on the scalar scatter rather than the popcount, so it only edges out SWAR by ~10%.
Make add_data use the AVX-512 method on x86-64 when the CPU supports avx512bitalg (runtime detected), falling back to SWAR otherwise.
Collaborator
|
Thank you for contribution! Notes for other reviewers: Running the bg4 predictor occupies only a tiny fraction of the upload runtime, so this is low-priority or not-planned. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add AVX 512 support for the bg4 predictor.
On my host with a Ryzen 9 7950X:
This is a modest increase in performance (10%) because the path is mostly bound on the scalar histogram scatter. Improving the scatter could be a good path forward for further performance increases.
Make the bg4 predictor also more generic over the word size, to allow using larger word sizes to fit 64-byte AVX 512 words, or possibly for other SIMD implementations later on. The AVX 512 implementation is runtime-detected, to allow making a single x86_64 build, and still get the performance increase where supported.
The
debug_asserts enforce the word-size contract documented onadd_data_implfor any future word type. The current u128 and __m512i instantiations always satisfy it. A compile-time check would needgeneric_const_exprs, which is still nightly-only.Note
Low Risk
Localized performance change to BG4 heuristic histograms with reference-equivalence tests and CPU feature fallbacks; no API or security-sensitive behavior changes beyond choosing a faster implementation of the same logic.
Overview
Adds a runtime-detected x86-64 fast path for BG4 histogram building using
vpopcntb(avx512bitalg), with SWAR fallback when the extension is missing.add_dataon x86-64 now routes through this path instead of SWAR-only.Refactors the inner driver to be generic over SIMD word size (
T: Podvia bytemuck), so the same alignment/head/body/tail loop works for 16-byte SWAR/Neon and 64-byte AVX-512, including 16-byte chunking for histogram scatter on wider words.Tests and the bg4_prediction benchmark now also exercise
add_data_avx512and assert histogram parity with the reference implementation.Reviewed by Cursor Bugbot for commit cde23c6. Bugbot is set up for automated code reviews on this repo. Configure here.